// ================================
// PCAUtils.cs
// PCA helper: mean centering, covariance, projection
// Path: Assets/LightProbePCA/Editor/PCAUtils.cs
// ================================

using System;
using UnityEditor;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;

public static class PCAUtils
{
    public enum PrincipalAxisMethod
    {
        Jacobi,
        SVD,
        CGS, //QR Classical Gram-Schmidt
        MGS  //QR Modified Gram-Schmidt
    }
    
    public static void CompressAndSave(int components, PrincipalAxisMethod method)
    {
#if  UNITY_EDITOR
        var probes = LightmapSettings.lightProbes;
        if (probes == null || probes.count == 0)
        {
            Debug.LogError("场景未烘培，没有light probes");
            return;
        }

        Vector3[] positions = probes.positions;
        SphericalHarmonicsL2[] shArray = probes.bakedProbes;

        // 转为矩阵
        float[][] M = PCAUtils.ToMatrix(shArray);

        // 计算均值、中心化、协方差
        float[] mean = PCAUtils.ComputeMean(M);
        float[][] C = PCAUtils.Center(M, mean);
        float[,] cov = PCAUtils.Covariance(C);

        float[,] axes = new float[components, cov.GetLength(1)]; 
        float[,] principalAxis = new float[components, cov.GetLength(1)]; 
        switch (method)
        {
            case PCAUtils.PrincipalAxisMethod.SVD:
            {
                // 使用 MathNet.Numerics SVD
                float[] singularValues;
                ComputeSVD(cov, out singularValues, out float[,] leftSingularVectors, out axes);
            }      
                break;
            case PCAUtils.PrincipalAxisMethod.Jacobi:
            {
                // Jacobi 特征分解获取主轴
                JacobiEigenDecomposition(cov, out float[] eigenValues, out axes);
            }
                break;
            case PCAUtils.PrincipalAxisMethod.CGS:
            {   // QR分解
                QRDecompose_CGS(cov, out axes, out float[,] upperTriangular);
            }
                break;
            case PCAUtils.PrincipalAxisMethod.MGS:
            {
                // QR分解
                QRDecompose_MGS(cov, out axes, out float[,] upperTriangular);
            }
                break;
        }

        // 取前 components 条主轴
        for (int i = 0; i < components; i++)
        for (int j = 0; j < cov.GetLength(1); j++)
            principalAxis[i, j] = axes[j, i]; // 使用 V 矩阵作为主轴     
        
        // 投影系数
        float[,] P = PCAUtils.Project(C, principalAxis);

        // 保存至 ScriptableObject
        string dir = "Assets/PCA/Data";
        if (!AssetDatabase.IsValidFolder(dir))
            AssetDatabase.CreateFolder("Assets/PCA", "Data");

        var asset = ScriptableObject.CreateInstance<CompressedProbeData>();
        asset.positions = positions;
        asset.mean = mean;
        asset.pcaAxes = new Float2DArray(principalAxis);
        asset.coefficients = new Float2DArray(P);
        
        string path = dir + "/CompressedProbeData.asset";
        AssetDatabase.CreateAsset(asset, path);
        AssetDatabase.SaveAssets();

        Debug.Log($"Compressed data saved to {path}");
#endif
    }

    public static CompressedProbeData CompressProbe(int components, PrincipalAxisMethod method)
    {
        var probes = LightmapSettings.lightProbes;
        if (probes == null || probes.count == 0)
        {
            Debug.LogError("场景未烘培，没有light probes");
            return null;
        }

        Vector3[] positions = probes.positions;
        SphericalHarmonicsL2[] shArray = probes.bakedProbes;

        // 转为矩阵
        float[][] M = PCAUtils.ToMatrix(shArray);

        // 计算均值、中心化、协方差
        float[] mean = PCAUtils.ComputeMean(M);
        float[][] C = PCAUtils.Center(M, mean);
        float[,] cov = PCAUtils.Covariance(C);

        float[,] axes = new float[components, cov.GetLength(1)]; 
        float[,] principalAxis = new float[components, cov.GetLength(1)]; 
        switch (method)
        {
            case PCAUtils.PrincipalAxisMethod.SVD:
            {
                // 使用 MathNet.Numerics SVD
                float[] singularValues;
                ComputeSVD(cov, out singularValues, out float[,] leftSingularVectors, out axes);
            } 
                break;
            case PCAUtils.PrincipalAxisMethod.Jacobi:
            {
                // Jacobi 特征分解获取主轴
                JacobiEigenDecomposition(cov, out float[] eigenValues, out axes);
            }
                break;
            case PCAUtils.PrincipalAxisMethod.CGS:
            {   // QR分解
                QRDecompose_CGS(cov, out axes, out float[,] upperTriangular);
            }
                break;
            case PCAUtils.PrincipalAxisMethod.MGS:
            {
                // QR分解
                QRDecompose_MGS(cov, out axes, out float[,] upperTriangular);
            }
                break;
        }

        // 取前 components 条主轴
        for (int i = 0; i < components; i++)
        for (int j = 0; j < cov.GetLength(1); j++)
            principalAxis[i, j] = axes[j, i]; // 使用 V 矩阵作为主轴
        
        // 投影系数
        float[,] P = PCAUtils.Project(C, principalAxis);

        CompressedProbeData compressedProbeData = ScriptableObject.CreateInstance<CompressedProbeData>();
        compressedProbeData.positions = positions;
        compressedProbeData.mean = mean;
        compressedProbeData.pcaAxes = new Float2DArray(principalAxis);
        compressedProbeData.coefficients = new Float2DArray(P);

        Debug.Log($"Compressed Data Success");
        
        return compressedProbeData;
    }
    public static void SaveData(CompressedProbeData compressedProbeData)
    {
        #if UNITY_EDITOR
        // 保存至 ScriptableObject
        string dir = "Assets/PCAGPU/Data";
        if (!AssetDatabase.IsValidFolder(dir))
            AssetDatabase.CreateFolder("Assets/PCAGPU", "Data");
        
        string path = dir + "/CompressedProbeData.asset";
        AssetDatabase.CreateAsset(compressedProbeData, path);
        AssetDatabase.SaveAssets();
        
        // 每个像素 3 通道，每通道 1 字节（UInt8）
        int size = 100;
        
        // 创建 Texture3D，使用 R8G8B8A8_UInt 格式
        var tex = new Texture3D(size, size, size, GraphicsFormat.R32_SFloat, TextureCreationFlags.None);
        tex.wrapMode = TextureWrapMode.Clamp;
        tex.filterMode = FilterMode.Point;

        Color[] pixels = new Color[size * size * size];
        for (int i = 0; i < pixels.Length; i++)
        {
            pixels[i] = new Color(1, 1, 1, 1);
        }
        tex.SetPixelData(pixels, 0);
        
        for (int i = 0;  i < compressedProbeData.positions.Length; i++)
        {
            tex.SetPixel((int)compressedProbeData.positions[i].x,(int)compressedProbeData.positions[i].y,(int)compressedProbeData.positions[i].z,Color.blue);
        }
        
        tex.Apply();
        
        path = dir + "/positionsRT.renderTexture";
        AssetDatabase.CreateAsset(tex, path);
        AssetDatabase.SaveAssets();

        Debug.Log($"Compressed data saved to {path}");
        
        //mean data lut
        // var tex_mean = new Texture2D(1, 9, GraphicsFormat.R8G8B8A8_SNorm, TextureCreationFlags.None);
        // tex_mean.wrapMode = TextureWrapMode.Clamp;
        // tex_mean.filterMode = FilterMode.Point;
        //
        // Color[] pixels_mean = new Color[9];
        // for (int i = 0; i < pixels_mean.Length; i++)
        // {
        //     pixels_mean[i] = new Color(1, 1, 1, 1);
        // }
        //
        // tex_mean.SetPixelData(pixels_mean, 0);
        //
        // for (int i = 0;  i < compressedProbeData.mean.Length;)
        // {
        //     tex_mean.SetPixel(0, i, new Color(compressedProbeData.mean[i + 0], compressedProbeData.mean[i + 1], compressedProbeData.mean[i + 2])); 
        //     i += 3;
        // }
        //
        // tex_mean.Apply();
        //
        // path = dir + "/mean.renderTexture";
        // AssetDatabase.CreateAsset(tex_mean, path);
        // AssetDatabase.SaveAssets();
        
        var tex_mean = new Texture2D(1, 27, GraphicsFormat.R32_SFloat, TextureCreationFlags.None);
        tex_mean.wrapMode = TextureWrapMode.Clamp;
        tex_mean.filterMode = FilterMode.Point;
        
        tex_mean.SetPixelData(compressedProbeData.mean, 0);
        
        tex_mean.Apply();
        
        path = dir + "/mean.renderTexture";
        AssetDatabase.CreateAsset(tex_mean, path);
        AssetDatabase.SaveAssets();
        
        //Axes data lut
        var tex_axes = new Texture2D(12, 27, GraphicsFormat.R32_SFloat, TextureCreationFlags.None);
        tex_axes.wrapMode = TextureWrapMode.Clamp;
        tex_axes.filterMode = FilterMode.Point;

        Color[] pixels_axes = new Color[12 * 27];
        for (int i = 0; i < pixels_axes.Length; i++)
        {
            pixels_axes[i] = new Color(1, 1, 1, 1);
        }
        
        tex_axes.SetPixelData(pixels_axes, 0);
        
        for (int i = 0;  i < compressedProbeData.pcaAxes.RowCount; i++)
        {
            for (int j = 0; j < compressedProbeData.pcaAxes.ColCount; j++)
            {
                tex_axes.SetPixel(i, j, new Color((compressedProbeData.pcaAxes[i, j]), 0, 0));
            }
        }
        
        tex_axes.Apply();
        
        path = dir + "/axes.renderTexture";
        AssetDatabase.CreateAsset(tex_axes, path);
        AssetDatabase.SaveAssets();
        
        //coefficients
        // 创建 Texture3D，使用 R8G8B8A8_UInt 格式
        var tex_coefficients = new Texture2D(compressedProbeData.coefficients.RowCount, compressedProbeData.coefficients.ColCount / 4, GraphicsFormat.R32G32B32A32_SFloat, TextureCreationFlags.None);
        tex_coefficients.wrapMode = TextureWrapMode.Clamp;
        tex_coefficients.filterMode = FilterMode.Point;

        // Color[] pixels_coefficients = new Color[compressedProbeData.coefficients.RowCount * compressedProbeData.coefficients.ColCount];
        // for (int i = 0; i < pixels_coefficients.Length; i++)
        // {
        //     pixels_coefficients[i] = new Color(1, 1, 1, 1);
        // }
        
        //tex_coefficients.SetPixelData(pixels, 0);
        
        for (int i = 0;  i < compressedProbeData.coefficients.RowCount; i++)
        {
            for (int j = 0; j < compressedProbeData.coefficients.ColCount;)
            {
                float x = compressedProbeData.coefficients[i, j + 0];
                float y = compressedProbeData.coefficients[i, j + 1];
                float z = compressedProbeData.coefficients[i, j + 2];
                float w = compressedProbeData.coefficients[i, j + 3];
                Vector4 v = new Vector4(x, y, z, w);
                tex_coefficients.SetPixel(i, j / 4, v);

                j += 4;
            }
        }
        
        tex_coefficients.Apply();
        
        path = dir + "/coefficients.renderTexture";
        AssetDatabase.CreateAsset(tex_coefficients, path);
        AssetDatabase.SaveAssets();
#endif
    }
    
    #region Decompose
    
    /// <summary>
    /// 使用 MathNet.Numerics 计算奇异值分解 
    /// </summary>
    /// <param name="samples">原始输入矩阵（m × n），每一行是一个样本，每一列是一个特征。</param>
    /// <param name="singularValues">长度为 min(m, n) 的数组，表示奇异值（等价于对称矩阵的特征值的平方根）。常用于表示主成分的能量或重要性。</param>
    /// <param name="leftSingularVectors">	矩阵 U（m × m）：每列是左奇异向量，表示样本空间中的正交基。</param>
    /// <param name="rightSingularVectors">矩阵 V（n × n）：每列是右奇异向量，也称为主轴，用于将原始特征投影到 PCA 空间。SVD 返回的其实是 V<sup>T</sup>，所以记得转置回来。</param>
    private static void ComputeSVD(float[,] samples, out float[] singularValues, out float[,] leftSingularVectors, out float[,] rightSingularVectors)
    {
        // 将二维数组转换为 MathNet 的矩阵类型
        var matrix = MathNet.Numerics.LinearAlgebra.Matrix<float>.Build.DenseOfArray(samples);

        // 计算 SVD
        var svd = matrix.Svd(true);

        // 提取 SVD 结果
        singularValues = svd.S.ToArray(); // 奇异值
        leftSingularVectors = svd.U.ToArray(); // U 矩阵
        rightSingularVectors = svd.VT.ToArray(); // V 矩阵（转置）

        // 注意：MathNet.Numerics 返回的 V 是 V.T，因此需要将其转置回来
    }
    
    /// <summary>
    /// Decompose symmetric matrix A into V (eigenvectors) and d (eigenvalues) 
    /// </summary>
    /// <param name="covarianceMatrix">输入的对称矩阵，通常是 协方差矩阵（n × n），必须是对称的。在 PCA 中，此矩阵是 X^T * X 或 (1/N) * X^T * X。</param>
    /// <param name="eigenValues">输出的特征值数组（长度 n），按照算法迭代过程中从对角线提取得到，未必排序，表示每个特征向量的“能量”或贡献度。</param>
    /// <param name="eigenVectors">输出的特征向量矩阵（n × n），每列是一个特征向量，通常被称为 主轴方向。这些向量是原始空间中最重要的“方向基”。</param>
    public static void JacobiEigenDecomposition(float[,] covarianceMatrix, out float[] eigenValues, out float[,] eigenVectors)
    {
        int n = covarianceMatrix.GetLength(0);
        eigenValues  = new float[n];
        eigenVectors  = new float[n, n];
        float[,] a = (float[,])covarianceMatrix.Clone();
        for (int i = 0; i < n; i++)
        {
            eigenValues [i] = a[i, i];
            for (int j = 0; j < n; j++)
                eigenVectors [i, j] = (i == j) ? 1f : 0f;
        }
    
        const int maxIter = 50;
        for (int iter = 0; iter < maxIter; iter++)
        {
            int p = 0, q = 1;
            float maxOff = Mathf.Abs(a[p, q]);
            for (int i = 0; i < n; i++)
            for (int j = i + 1; j < n; j++)
            {
                float val = Mathf.Abs(a[i, j]);
                if (val > maxOff)
                {
                    maxOff = val;
                    p = i;
                    q = j;
                }
            }
    
            if (maxOff < 1e-6f) break;
            float theta = 0.5f * (a[q, q] - a[p, p]) / a[p, q];
            float t = 1f / (Math.Abs(theta) + Mathf.Sqrt(1f + theta * theta));
            if (theta < 0f) t = -t;
            float c = 1f / Mathf.Sqrt(1f + t * t);
            float s = t * c;
            for (int i = 0; i < n; i++)
            {
                if (i == p || i == q) continue;
                float aip = a[i, p], aiq = a[i, q];
                a[i, p] = a[p, i] = c * aip - s * aiq;
                a[i, q] = a[q, i] = c * aiq + s * aip;
            }
    
            float app = a[p, p], aqq = a[q, q], apq = a[p, q];
            a[p, p] = c * c * app - 2f * s * c * apq + s * s * aqq;
            a[q, q] = s * s * app + 2f * s * c * apq + c * c * aqq;
            a[p, q] = a[q, p] = 0f;
            for (int i = 0; i < n; i++)
            {
                float vip = eigenVectors [i, p], viq = eigenVectors [i, q];
                eigenVectors [i, p] = c * vip - s * viq;
                eigenVectors [i, q] = s * vip + c * viq;
            }
        }
    
        for (int i = 0; i < n; i++) eigenValues [i] = a[i, i];
    }
    
    
    /// <summary>
    /// QR Classical Gram-Schmidt, CGS 分解的实现 
    /// </summary>
    /// <param name="samples">矩阵 A（m × n），每列是一个向量（用于构造正交基）</param>
    /// <param name="orthonormalBasis">	Q 矩阵（m × n），正交单位列向量组，构成输入向量的正交基。即：列向量两两正交且单位长度。</param>
    /// <param name="upperTriangular">R 矩阵（n × n），上三角矩阵，记录了原始向量在正交基下的坐标变换系数（投影系数）</param>
    private static void QRDecompose_CGS(float[,] samples, out float[,] orthonormalBasis, out float[,] upperTriangular)
    {
        int rows = samples.GetLength(0);
        int cols = samples.GetLength(1);
        orthonormalBasis = new float[rows, cols];
        upperTriangular = new float[cols, cols];

        // 将 A 的列复制到列向量
        float[][] aCols = new float[cols][];
        for (int j = 0; j < cols; j++)
            aCols[j] = GetColumn(samples, j);

        float[][] qCols = new float[cols][];

        for (int k = 0; k < cols; k++)
        {
            float[] vk = (float[])aCols[k].Clone();

            // 逐列减去投影部分
            for (int j = 0; j < k; j++)
            {
                float r = Dot(qCols[j], aCols[k]);
                upperTriangular[j, k] = r;
                for (int i = 0; i < rows; i++)
                    vk[i] -= r * qCols[j][i];
            }

            float norm = Norm(vk);
            upperTriangular[k, k] = norm;

            float[] qk = new float[rows];
            for (int i = 0; i < rows; i++)
                qk[i] = vk[i] / norm;

            qCols[k] = qk;
        }

        // 将 qCols 转换为 Q
        for (int j = 0; j < cols; j++)
            for (int i = 0; i < rows; i++)
                orthonormalBasis[i, j] = qCols[j][i];
    }

    
    /// <summary>
    ///QR Modified Gram-Schmidt, MGS 
    /// </summary>
    /// <param name="samples">矩阵 A（m × n），每列是一个向量（用于构造正交基）</param>
    /// <param name="orthonormalBasis">	Q 矩阵（m × n），正交单位列向量组，构成输入向量的正交基。即：列向量两两正交且单位长度。</param>
    /// <param name="upperTriangular">R 矩阵（n × n），上三角矩阵，记录了原始向量在正交基下的坐标变换系数（投影系数）</param>
    private static void QRDecompose_MGS(float[,] samples, out float[,] orthonormalBasis, out float[,] upperTriangular)
    {
        int rows = samples.GetLength(0);
        int cols = samples.GetLength(1);
        orthonormalBasis = new float[rows, cols];
        upperTriangular = new float[cols, cols];

        // 取列向量
        float[][] qCols = new float[cols][];
        for (int k = 0; k < cols; k++)
        {
            // 复制 A 的第 k 列作为初始向量
            float[] qk = GetColumn(samples, k);
            for (int j = 0; j < k; j++)
            {
                float r = Dot(qCols[j], qk);
                upperTriangular[j, k] = r;
                for (int i = 0; i < rows; i++)
                    qk[i] -= r * qCols[j][i];
            }

            float norm = Norm(qk);
            upperTriangular[k, k] = norm;

            for (int i = 0; i < rows; i++)
                qk[i] /= norm;

            qCols[k] = qk;
        }

        // 转回 Q 矩阵
        for (int j = 0; j < cols; j++)
        for (int i = 0; i < rows; i++)
            orthonormalBasis[i, j] = qCols[j][i];
    }

    // 获取矩阵 A 的某一列
    private static float[] GetColumn(float[,] matrix, int col)
    {
        int rows = matrix.GetLength(0);
        float[] column = new float[rows];
        for (int i = 0; i < rows; i++)
            column[i] = matrix[i, col];
        return column;
    }
    
    #endregion
    
    private static float[][] ToMatrix(SphericalHarmonicsL2[] sh)
    {
        int n = sh.Length;
        float[][] M = new float[n][];
        for (int i = 0; i < n; i++)
        {
            M[i] = new float[27];
            for (int c = 0; c < 3; c++)
            for (int j = 0; j < 9; j++)
                M[i][c * 9 + j] = sh[i][c, j];
        }

        return M;
    }

    private static float[] ComputeMean(float[][] M)
    {
        int n = M.Length;
        int d = M[0].Length;
        float[] mean = new float[d];
        for (int i = 0; i < n; i++)
        for (int j = 0; j < d; j++)
            mean[j] += M[i][j];
        for (int j = 0; j < d; j++)
            mean[j] /= n;
        return mean;
    }

    private static float[][] Center(float[][] M, float[] mean)
    {
        int n = M.Length;
        int d = mean.Length;
        float[][] C = new float[n][];
        for (int i = 0; i < n; i++)
        {
            C[i] = new float[d];
            for (int j = 0; j < d; j++)
                C[i][j] = M[i][j] - mean[j];
        }

        return C;
    }

    private static float[,] Covariance(float[][] C)
    {
        int n = C.Length;
        int d = C[0].Length;
        float[,] cov = new float[d, d];
        for (int i = 0; i < n; i++)
        for (int a = 0; a < d; a++)
        for (int b = 0; b < d; b++)
            cov[a, b] += C[i][a] * C[i][b];
        float inv = 1f / (n - 1);
        for (int a = 0; a < d; a++)
        for (int b = 0; b < d; b++)
            cov[a, b] *= inv;
        return cov;
    }

    // 返回投影系数矩阵：P[n_samples, n_components]
    private static float[,] Project(float[][] C, float[,] axes)
    {
        int n = C.Length;
        int k = axes.GetLength(0);
        int d = axes.GetLength(1);
        float[,] P = new float[n,k];
        for (int i = 0; i < n; i++)
        {
            //P[i] = new float[k];
            for (int a = 0; a < k; a++)
            {
                float sum = 0f;
                for (int j = 0; j < d; j++)
                    sum += C[i][j] * axes[a, j];
                P[i,a] = sum;
            }
        }

        return P;
    }
    
    // 向量点积
    private static float Dot(float[] a, float[] b)
    {
        float sum = 0f;
        for (int i = 0; i < a.Length; i++)
            sum += a[i] * b[i];
        return sum;
    }

    // 向量的范数（即 L2 范数）
    private static float Norm(float[] v)
    {
        return Mathf.Sqrt(Dot(v, v));
    }

    // 矩阵乘法
    private static float[,] Multiply(float[,] A, float[,] B)
    {
        int rowsA = A.GetLength(0);
        int colsA = A.GetLength(1);
        int colsB = B.GetLength(1);
        float[,] result = new float[rowsA, colsB];

        for (int i = 0; i < rowsA; i++)
        for (int j = 0; j < colsB; j++)
        for (int k = 0; k < colsA; k++)
            result[i, j] += A[i, k] * B[k, j];

        return result;
    }

    // 创建单位矩阵
    private static float[,] Identity(int size)
    {
        float[,] id = new float[size, size];
        for (int i = 0; i < size; i++)
            id[i, i] = 1f;
        return id;
    }    
}